"""
@author: 景云鹏
@email: 310491287@qq.com
@date: 2022/5/18
"""
import time
from abc import ABCMeta, abstractmethod
from contextlib import contextmanager

from numpy import *


class DirectionProvider(metaclass=ABCMeta):

    @abstractmethod
    def get_dk(self, context):
        raise NotImplementedError


class StepProvider(metaclass=ABCMeta):
    @abstractmethod
    def get_alpha(self, context, dk):
        raise NotImplementedError


class Context:
    def __init__(self, f, g, x0):
        self.f = f
        self.g = g
        self.opt = None
        self.records = []
        self.xk = x0

    @property
    def xk(self):
        return self.records[-1][0]

    @property
    def fk(self):
        return self.records[-1][1]

    @property
    def gk(self):
        return self.records[-1][2]

    @xk.setter
    def xk(self, x_value):
        fk = self.f(x_value)
        if self.opt is None or self.opt > fk:
            self.opt = fk
        self.records.append((x_value, fk, self.g(x_value)))

    @contextmanager
    def analysis(self):
        t1 = time.time()
        yield
        iters = len(self.records)
        cpu = time.time() - t1
        calls = self.f.counter + self.g.counter
        print(dict(
            iters=iters, cpu=cpu, calls=calls, opt=self.opt
        ))


class Count:
    def __init__(self, func):
        self.func = func
        self.counter = 0

    def __call__(self, *args, **kwargs):
        self.counter += 1
        return self.func(*args, **kwargs)


class LinearOptimizer:
    def __init__(self, dp: DirectionProvider, sp: StepProvider):
        self.dp = dp
        self.sp = sp

    def optimize(self, c: Context):
        dk = self.dp.get_dk(c)
        while abs(c.gk).max() > 1e-5:
            try:
                ak = self.sp.get_alpha(c, dk)
            except ValueError:
                break
            c.xk = c.xk + ak * dk
            dk = self.dp.get_dk(c)
        return c.fk


class Armijo(StepProvider):
    def __init__(self, beta=0.5, sigma=0.2, max_k=20):
        self.beta = beta
        self.sigma = sigma
        self.max_k = max_k

    def get_alpha(self, context, dk):
        mk = 0
        bm = 1
        xk = context.xk
        f, fk, gk = context.f, context.fk, context.gk
        sgd = self.sigma * dot(gk, dk)
        while mk <= self.max_k:
            if f(xk + bm * dk) <= fk + bm * sgd:
                return bm
            bm *= self.beta
            mk += 1
        return bm


class SteepestDescent(DirectionProvider):

    def get_dk(self, context):
        return -context.gk


class CG(DirectionProvider, metaclass=ABCMeta):
    def __init__(self):
        self.last_dk = None
        self.last_gk = None

    def get_dk(self, context):
        if len(context.records) == 1:
            dk = -context.gk
        else:
            beta = self.get_beta(context)
            dk = beta * self.last_dk - context.gk
        self.last_gk = context.gk
        self.last_dk = dk
        return dk

    @abstractmethod
    def get_beta(self, context):
        raise NotImplementedError


class FRCG(CG):
    def get_beta(self, context):
        return dot(context.gk, context.gk) / dot(
            self.last_gk, self.last_gk
        )


class PRPCG(CG):
    def get_beta(self, context):
        beta = dot(context.gk, context.gk - self.last_gk) / dot(
            self.last_gk, self.last_gk
        )
        return max((beta, 0))


class HSCG(CG):
    def get_beta(self, context):
        dif = context.gk - self.last_gk
        return dot(context.gk, dif) / dot(dif, self.last_dk)


class DYCG(CG):
    def get_beta(self, context):
        return dot(context.gk, context.gk) / dot(
            context.gk - self.last_gk, self.last_dk
        )


class BFGS(DirectionProvider):
    def __init__(self):
        self.last_bk = None

    def get_dk(self, context):
        n = len(context.xk)
        if len(context.records) == 1:
            bk = eye(n)
        else:
            sk = context.xk - context.records[-2][0]
            yk = context.gk - context.records[-2][2]
            sk = mat(sk.reshape((n, 1)))
            yk = mat(yk.reshape((n, 1)))
            second = self.last_bk * sk * sk.T * self.last_bk / (sk.T * self.last_bk * sk)
            third = yk * yk.T / (sk.T * sk)
            bk = self.last_bk - second + third
        self.last_bk = bk
        result = matmul(-linalg.inv(bk), context.gk)
        if isinstance(result, matrix):
            result = result.A.reshape(4)
        return result


@Count
def f(x):
    xm = x - 1
    return sum((
        100 * (x[0] ** 2 - x[1]) ** 2,
        xm[0] ** 2 + xm[2] ** 2,
        90 * (x[2] ** 2 - x[3]) ** 2,
        10.1 * (xm[1] ** 2 + xm[3] ** 2),
        19.8 * xm[1] * xm[3] + 10
    ))


@Count
def g(x):
    xm = x - 1
    return asarray((
        400 * x[0] * (x[0] ** 2 - x[1]) + 2 * xm[0],
        -200 * (x[0] ** 2 - x[1]) + 20.2 * xm[1] + 19.8 * xm[3],
        360 * x[2] * (x[2] ** 2 - x[3]) + 2 * xm[2],
        -180 * (x[2] ** 2 - x[3]) + 20.2 * xm[3] + 19.8 * xm[1]
    ))


def run(x0, dp, sp):
    lin = LinearOptimizer(dp, sp)
    c = Context(f, g, x0)
    with c.analysis():
        lin.optimize(c)


def get_x0_list():
    xes = []
    for i in [0.1, 1, 10]:
        for j in [asarray((-1, 1, -1, -1)), asarray((-5, 5, -5, 5))]:
            xes.append(i * j)
    return xes


def get_direction_providers():
    return [
        SteepestDescent(),
        PRPCG(),
        HSCG(),
        DYCG(),
        FRCG(),
        BFGS()
    ]


if __name__ == '__main__':
    import warnings

    warnings.filterwarnings('ignore')
    arm = Armijo()
    for x0 in get_x0_list():
        for dp in get_direction_providers():
            print(x0, dp.__class__.__name__)
            run(x0, dp, arm)
            print('========================')
            print()
